import torch  # 数据预处理
from torch import nn  # 模型设计

# 数据集的二次加工工具
from torchvision import transforms
# 加载图片数据集，并且自动分类
from torchvision.datasets import ImageFolder
from NumberNet import NumberNet

# -----------------加载数据集-------------------
root_dir = "./data/numbers"
# 1--对数据集进行二次加工
transform = transforms.Compose([
    transforms.ToTensor()  # 希望加载的图片数据为pytorch的tensor对象
])
# 2--加载所有的图片数据，并且根据文件夹结构分类
number_datasets = ImageFolder(root_dir, transform=transform)
# 3--为数据设置batch_size,以及shuffle
#    batch_size:定义每次训练需要图片数量， shuffle：将原有的数据打乱顺序（训练阶段）
train_loader = torch.utils.data.DataLoader(number_datasets, batch_size=1000, shuffle=True)


# 验证loader中的数据
# images, labels = next(iter(train_loader))
# print(images.shape)
# print(labels.shape)

# -----------------模型设计-------------------
model = NumberNet(classes=10)
# -----------------损失函数-------------------
criterion = nn.CrossEntropyLoss()  # 交叉熵 + softmax
# -----------------优化器-------------------
sgd = torch.optim.SGD(model.parameters(), lr=0.1)
# -----------------训练图像-----------------
epochs = 10
for epoch in range(epochs):
    print(f"----epoch {epoch + 1} / {epochs}:")
    for images, labels in iter(train_loader):
        sgd.zero_grad()
        predict_labels = model(images.to(model.device))
        loss = criterion(predict_labels, labels.to(model.device))
        loss.backward()
        sgd.step()

        print(f"------loss:{loss.item():.4f}")

# 保存最终的模型
torch.save(model, "numbers_model2.pth")
# 优化的过程：
# 1、针对过多的神经元进行削减（只针对nn.Linear）
# 2、Softmax变为LogSoftmax 进行熵化结果
# 3、调整学习率、调整训练次数...
